#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 19-2-12 下午5:40
# @Author  : Jh Zhao
# @Site    : 
# @File    : annotation_utils.py
# @Software: PyCharm Community Edition

from lxml import etree
import numpy as np


def _recursive_parse_xml_to_dict(xml):
    """Recursively parses XML contents to python dict.

    We assume that `object` and `point` tags are the only twos that can appear
    multiple times at the same level of a tree.

    Args:
      xml: xml tree obtained by parsing XML file contents using lxml.etree

    Returns:
      Python dictionary holding XML contents.
    """
    if xml is None:
        return {}
    if len(xml) == 0:
        return {xml.tag: xml.text}
    result = {}
    for child in xml:
        child_result = _recursive_parse_xml_to_dict(child)
        if child.tag not in ('object', 'point'):
            result[child.tag] = child_result[child.tag]
        else:
            if child.tag not in result:
                result[child.tag] = []
            result[child.tag].append(child_result[child.tag])
    return {xml.tag: result}


def _recursive_create_dict_to_xml(dct, root):
    """Recursively create XML contents base on a python dict.

    Args:
      dct: python dictionary holding XML contents
      root: xml tree root where dict contents will append on.

    Returns:
    """
    for key, val in dct.items():
        if isinstance(val, dict):
            node = etree.SubElement(root, key)
            _recursive_create_dict_to_xml(val, node)
        elif isinstance(val, list):
            for x in val:
                node = etree.SubElement(root, key)
                _recursive_create_dict_to_xml(x, node)
        else:
            node = etree.SubElement(root, key)
            node.text = str(val)


def read_annotation(annotation_path):
    """ Read object detection annotation of xml format file
    
    Args:
      annotation_path: file path of annotation

    Returns: a dict of detection annotation
    """
    with open(annotation_path, 'rb') as f:
        xml_str = f.read()
    xml = etree.fromstring(xml_str)
    data = _recursive_parse_xml_to_dict(xml)["annotation"]
    return data


def write_annotation(annotation_path, annotation):
    """ Write object detection annotation to a xml format file
    
    Args:
      annotation_path: file path of annotation
      annotation: a dict of detection annotation

    Returns: 
    """
    root = etree.Element("annotation")
    _recursive_create_dict_to_xml(annotation, root)  # write to file:
    tree = etree.ElementTree(root)
    tree.write(annotation_path, pretty_print=True, encoding='utf-8')


def _round_int(x):
    return int(round(x))


def get_polygon_from_obj(obj):
    """ Get polygon from object annotation

    :param obj: objects's annotation
    :return: a np.ndarray with shape [N, 2] as a list of points consists the polygon.
        A point is arranged as (x, y)
    """
    if "polygon" not in obj:
        raise ValueError("must have a polygon in object")
    pts = [[_round_int(float(p['x'])), _round_int(float(p['y']))]
           for p in obj['polygon'].get('point', [])]
    return np.asarray(pts)


def set_polygon_to_obj(obj, polygon):
    """ Set polygon to object annotation

    :param obj: objects's annotation
    :param polygon: a list of points [(x0, y0), ...]
    :return:
    """
    points = [{"x": x, "y": y} for x, y in polygon]
    polygon = {"point": points}
    obj["polygon"] = polygon


def get_bndbox_from_obj(obj):
    """ Get bounding box from object annotation

    :param obj: objects's annotation
    :return: a tuple consists of ymin, xmin, ymax, xmax
    """
    if "bndbox" not in obj:
        raise ValueError("must have a bndbox in object")
    bndbox = [_round_int(float(obj["bndbox"][x])) for x
              in ["ymin", "xmin", "ymax", "xmax"]]
    return tuple(bndbox)


def set_bndbox_to_obj(obj, bndbox):
    """ Set bounding box to object annotation

    :param obj: objects's annotation
    :param bndbox: a tuple or list consists of ymin, xmin, ymax, xmax
    :return:
    """
    bndbox = {"ymin": bndbox[0], "xmin": bndbox[1],
              "ymax": bndbox[2], "xmax": bndbox[3]}
    obj["bndbox"] = bndbox


def get_polygons_from_ann(ann):
    """ Get all polygons from an annotation

    :param ann: objects's annotation
    :return: a list of polygon in annotation
    """
    objs = ann.get('object', [])
    return [get_polygon_from_obj(x) for x in objs]


def set_polygons_to_ann(ann, polygons):
    """ Set all polygons to an annotation

    :param ann: objects's annotation
    :param polygons: a list of polygons
    :return:
    """
    objs = ann.get('object', [])
    assert len(objs) == polygons, \
        "the length of polygons must be same as objects"
    list(map(set_polygon_to_obj, objs, polygons))


def get_bndboxes_from_ann(ann):
    """ Get all bounding boxes from an annotation

    :param ann: objects's annotation
    :return: a list of bnd boxes in annotation
    """
    objs = ann.get('object', [])
    return [get_bndbox_from_obj(x) for x in objs]


def set_bndboxes_to_ann(ann, bndboxes):
    """ Set all bounding boxes to an annotation

    :param ann: objects's annotation
    :return: a list of bnd boxes in annotation
    """
    objs = ann.get('object', [])
    assert len(objs) == bndboxes, \
        "the length of bndboxes must be same as objects"
    list(map(set_bndbox_to_obj, objs, bndboxes))


def extract_bndbox_from_object(obj):
    """ Extract a bounding box from object annotation

        If there is a bndbox in object annotation, then get it. If there is
        a polygon in there, get it and use its minimum bounding box as result.
        Otherwise, raise a ValueError.

    :param obj: objects's annotation
    :return: a tuple consists of ymin, xmin, ymax, xmax
    """
    if "bndbox" in obj:
        return get_bndbox_from_obj(obj)
    elif "polygon" in obj:
        plg = get_polygon_from_obj(obj)
        xmin, ymin = np.min(plg, axis=0)
        xmax, ymax = np.max(plg, axis=0)
        return ymin, xmin, ymax, xmax
    else:
        raise ValueError("not 'bndbox' or 'polygon' in obj")


def extract_groundtruth_from_ann(ann, ignore_obj_fn=None,
                                 ignore_reg_fn=None):
    """ Extract ground-truth from an annotation

    :param ann: a dict of annotation information
    :param ignore_obj_fn: a callable function receive object annotation
        and return bool which determines ignore one object or not.
    :param ignore_reg_fn: a callable function receive object annotation
        and return bool which determines whether regard one object as a
        ignore region or not.
    :return: a dict contains ground-truth from a annotation
        "boxes": a np.ndarray with shape [N, 4]
        "classes": a np.ndarray with shape [N]
        "ignore_boxes": optional, a np.ndarray with shape [N]
        "ignore_regions": optional, a np.ndarray with shape [N]
    """
    objs = ann.get('object', [])
    g_bxs = [extract_bndbox_from_object(x) for x in objs]
    g_cls = [x['name'] for x in objs]
    g_res = {"boxes": np.asarray(g_bxs).reshape([-1, 4]),
             "classes": np.asarray(g_cls)}
    if ignore_obj_fn:
        g_igr = [ignore_obj_fn(obj) for obj in objs]
        g_res["ignore_boxes"] = np.asarray(g_igr, dtype=np.bool)
    if ignore_reg_fn:
        g_igr_r = [ignore_reg_fn(obj) for obj in objs]
        g_res["ignore_regions"] = np.asarray(g_igr_r, dtype=np.bool)
    return g_res


def extract_predict_from_ann(ann, ignore_obj_fn=None):
    """ Extract predict result from an annotation

    :param ann: a dict of annotation information
    :param ignore_obj_fn: a callable function receive object annotation
        and return bool which determines ignore one object or not.
    :return: a dict contains ground-truth from a annotation
        "boxes": a np.ndarray with shape [N, 4]
        "classes": a np.ndarray with shape [N]
        "scores": a np.ndarray with shape [N]
        "ignore_boxes": optional, a np.ndarray with shape [N]
    """
    objs = ann.get('object', [])
    g_bxs = [extract_bndbox_from_object(x) for x in objs]
    g_cls = [x['name'] for x in objs]
    g_scr = [float(x.get('score', '1')) for x in objs]
    g_res = {"boxes": np.asarray(g_bxs).reshape([-1, 4]),
             "classes": np.asarray(g_cls),
             "scores": np.asarray(g_scr)}
    if ignore_obj_fn:
        g_igr = [ignore_obj_fn(obj) for obj in objs]
        g_res["ignore_boxes"] = np.asarray(g_igr, dtype=np.bool)
    return g_res


def create_object(name, *, bndbox=None, polygon=None, status=None):
    """ Create detection object

    object should set one bndbox or polygon

    :param name: object name
    :param bndbox: bounding box as [ymin, xmin, ymax, xmax]
    :param polygon: polygon as [(x0, y0), (x1, y1), ...]
    :param status: one integer indicates label status
    :return: a dict represents a detection object
    """
    assert (bndbox or polygon) and not (bndbox and polygon), \
        "One of bndbox or polygon must be set"
    obj = {"name": name}
    if status:
        obj['status'] = status
    if bndbox:
        obj['bndbox'] = {"ymin": bndbox[0], "xmin": bndbox[1],
                         "ymax": bndbox[2], "xmax": bndbox[3]}
    else:
        obj['polygon'] = {"point": [{"x": x, "y": y} for x, y in polygon]}
    return obj


def create_annotation(filename=None, url=None, size=None, segmented=None,
                      objects=None):
    """ Create detection annotation

    :param filename: image file name
    :param url: image file url
    :param size: image size, like (width, height) or (width ,height, depth)
    :param segmented:
    :param objects: a list of objects
    :return: a dict represents a detection annotation
    """
    ann = {}
    if filename:
        ann['filename'] = filename
    if url:
        ann['url'] = url
    if size:
        ann['size'] = {k: v for k, v in zip(
            ("width", "height", "depth"), size)}
    if segmented:
        ann['segmented'] = segmented
    if objects:
        ann['object'] = objects
    return ann
